from django.shortcuts import render
from django.views.generic import View
from scan.component import fengxian, geo, saomiao

from django.http import HttpResponse
import json
import os
import pandas as pd
from DecisionTree.test import DecisionTree_train
from getData.datadb import sql_ssvid
from scan.pocc import pocc


# Create your views here.
class ScanView(View):

    def get(self, request):
        return render(request, "index.html", {})

    def post(self, request):
        keyword = request.POST.get("input", "")
        keyword1 = request.POST.get("way", "")

        # 全面扫描
        if keyword1 != 'Quick':
            way = "-PE -n --min-hostgroup 1024 --min-parallelism 1024 -F -T4 -Pn -sS -v -O"
        else:
            # 快速扫描
            ip_str, ip_port_str = self._faster(keyword)
            if ip_str.replace(" ", "") == "" or ip_port_str.replace(" ", "") == "":
                way = "-PE -n --min-hostgroup 1024 --min-parallelism 1024 -T4 -Pn -sS -A"
            else:
                way = F" {ip_str} -PE -n --min-hostgroup 1024 --min-parallelism 1024 -T4  -sS -v -O -p {ip_port_str}"
            print(way)
        output = saomiao.saomiao(str(keyword), way)
        print(output)
        # 根据决策树匹配cve
        BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        data = pd.read_csv(os.path.join(BASE_DIR, "DecisionTree/os_vul.csv"))
        dt = DecisionTree_train(data)
        cve = dt.predict_result(dt.get_inf(output))
        print("=============\n", cve)

        # 把cve转换为ssvid
        ssvid = sql_ssvid().get_ssvid(cve)
        print("=============\n", ssvid)

        # 扫描poc
        ip = [x["host"] for x in output]
        poc_result = pocc.pocc().get_result(ip, ssvid)
        print("=============\n", poc_result)

        geo_ip = geo.geo(output)
        risk = fengxian.fengxian(output)
        r = {"result": output, "geo_ip1": geo_ip, "fengxian1": risk}
        print(r)
        return HttpResponse(json.dumps(r))
        # return render_template('../../../python/Network_Identify/templates/index1.html',result = output,geo_ip1 = geo_ip,fengxian1 = risk)
        # return ujson.dumps(output)

    def _faster(self, keyword):
        # 使用nmap扫描得到存活主机
        way = "-sP"
        output = saomiao.saomiao(str(keyword), way)
        # tmpfile = tempfile.mktemp()
        # tmpfile_2 = tempfile.mktemp()
        tmpfile = '/tmp/live_ip_result.list'
        tmpfile_2 = '/tmp/live_ip.list'
        with open(tmpfile_2, 'w') as file:
            for i in range(0, len(output)):
                live_ip = output[i]["host"]
                file.write(live_ip + '\n')
        # 使用masscan扫描存活主机
        exe = F'masscan -iL {tmpfile_2} -p 1-65535 --rate 20000 -oL {tmpfile} --wait 0'
        # exe = f'sudo masscan -iL /tmp/live_ip.list -p 1-65535 --rate 20000 -oL /tmp/1.list --wait 0'

        os.system(exe)

        # 进行massan结果读取
        with open(tmpfile, "r") as file:
            ip_port_dict = {}  # 定义字典
            # 将结果文件转为ip：port的字典
            for line_data in file:
                line_list = line_data.split()
                if line_list[0] == 'open':
                    ip = line_list[3]
                    port = line_list[2]
                    if ip not in ip_port_dict:
                        ip_port_dict[ip] = []
                    ip_port_dict[ip].append(port)
        ip_addresses = []
        ip_ports = []
        for ip, port in ip_port_dict.items():
            ip_addresses.append(ip)
            ip_ports.extend(port)
        ip_ports = set(ip_ports)
        ip_addresses_str = ""
        ip_ports_str = ""
        for i in ip_addresses:
            ip_addresses_str += str(i + ' ')
        for i in ip_ports:
            ip_ports_str += str(i + ' ')
        return ip_addresses_str, ip_ports_str


class getCVEView(View):

    def post(self, request):
        keyword = request.POST.get("ip", "")
        method = request.POST.get("method", "")
        # 全面扫描
        if method != 'Quick':
            way = "-PE -n --min-hostgroup 1024 --min-parallelism 1024 -F -T4 -Pn -sS -v -O"
        else:
            # 快速扫描
            ip_str, ip_port_str = self._faster(keyword)
            if ip_str.replace(" ", "") == "" or ip_port_str.replace(" ", "") == "":
                way = "-PE -n --min-hostgroup 1024 --min-parallelism 1024 -T4 -Pn -sS -A"
            else:
                way = F" {ip_str} -PE -n --min-hostgroup 1024 --min-parallelism 1024 -T4  -sS -v -O -p {ip_port_str}"
            print(way)

        output = saomiao.saomiao(str(keyword), way)
        # 根据决策树匹配cve
        BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

        data = pd.read_csv(os.path.join(BASE_DIR, "DecisionTree/os_vul.csv"))
        dt = DecisionTree_train(data)
        cve = dt.predict_result(dt.get_inf(output))
        print(cve)
        # 把cve转换为ssvid
        ssvid = sql_ssvid().get_ssvid(cve)
        print(ssvid)
        if ssvid == [None]:
            ssvid = None
            ssvid_url = None
        else:
            # 扫描poc
            ip = [x["host"] for x in output]
            poc_result = pocc.pocc().get_result(ip, ssvid[0])
            print("=============\n", poc_result)

            ssvid = ssvid[0].split('-')[-1]
            ssvid_url = "https://www.seebug.org/vuldb/ssvid-" + ssvid

        r = {"cve": cve[0], "ssvid": ssvid, "url": ssvid_url}
        print(r)
        return HttpResponse(json.dumps(r))

    def _faster(self, keyword):
        # 使用nmap扫描得到存活主机
        way = "-sP"
        output = saomiao.saomiao(str(keyword), way)
        # tmpfile = tempfile.mktemp()
        # tmpfile_2 = tempfile.mktemp()
        tmpfile = '/tmp/live_ip_result.list'
        tmpfile_2 = '/tmp/live_ip.list'
        with open(tmpfile_2, 'w') as file:
            for i in range(0, len(output)):
                live_ip = output[i]["host"]
                file.write(live_ip + '\n')
        # 使用masscan扫描存活主机
        exe = F'masscan -iL {tmpfile_2} -p 1-65535 --rate 20000 -oL {tmpfile} --wait 0'
        # exe = f'sudo masscan -iL /tmp/live_ip.list -p 1-65535 --rate 20000 -oL /tmp/1.list --wait 0'

        os.system(exe)

        # 进行massan结果读取
        with open(tmpfile, "r") as file:
            ip_port_dict = {}  # 定义字典
            # 将结果文件转为ip：port的字典
            for line_data in file:
                line_list = line_data.split()
                if line_list[0] == 'open':
                    ip = line_list[3]
                    port = line_list[2]
                    if ip not in ip_port_dict:
                        ip_port_dict[ip] = []
                    ip_port_dict[ip].append(port)
        ip_addresses = []
        ip_ports = []
        for ip, port in ip_port_dict.items():
            ip_addresses.append(ip)
            ip_ports.extend(port)
        ip_ports = set(ip_ports)
        ip_addresses_str = ""
        ip_ports_str = ""
        for i in ip_addresses:
            ip_addresses_str += str(i + ' ')
        for i in ip_ports:
            ip_ports_str += str(i + ' ')
        return ip_addresses_str, ip_ports_str
